import cv2
import os
import numpy as np

# 配置参数
input_folder = "D:/Desk/test1/content/"  # 原始图片目录
output_folder = "D:/Desk/test1/result/"  # 输出目录
target_ratio = 3/4  # 宽高比 3:4

def detect_clothing_center(img):
    """通过颜色和纹理分析定位衣物中心"""
    h, w = img.shape[:2]
    
    # --- 步骤1：人体区域粗略检测 ---
    hog = cv2.HOGDescriptor()
    hog.setSVMDetector(cv2.HOGDescriptor_getDefaultPeopleDetector())
    (regions, _) = hog.detectMultiScale(img, winStride=(4,4), padding=(8,8), scale=1.05)
    
    if len(regions) > 0:
        # 取面积最大的人体区域
        x, y, rw, rh = max(regions, key=lambda r: r[2]*r[3])
        # 聚焦于人体上半身（假设衣物区域占检测框的60%高度）
        roi = img[y:y+int(rh*0.6), x:x+rw]
    else:
        roi = img  # 无检测时使用全图
    
    # --- 步骤2：颜色聚类定位主衣物区域 ---
    hsv = cv2.cvtColor(roi, cv2.COLOR_BGR2HSV)
    pixels = hsv.reshape((-1, 3))
    pixels = np.float32(pixels)
    
    # K-means聚类（寻找主要颜色）
    criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 10, 1.0)
    _, labels, centers = cv2.kmeans(pixels, 2, None, criteria, 10, cv2.KMEANS_RANDOM_CENTERS)
    
    # 选择面积最大的颜色区域
    dominant_color = centers[np.argmax(np.bincount(labels.flatten()))]
    
    # 创建颜色掩膜
    lower = np.clip(dominant_color - [10, 50, 50], 0, 255)
    upper = np.clip(dominant_color + [10, 80, 80], 0, 255)
    mask = cv2.inRange(hsv, lower, upper)
    
    # --- 步骤3：计算衣物几何中心 ---
    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    if contours:
        # 合并所有轮廓
        all_contours = np.vstack(contours)
        # 计算最小包围矩形
        rect = cv2.boundingRect(all_contours)
        # 转换为全局坐标
        if roi is not img:
            rect = (rect[0]+x, rect[1]+y, rect[2], rect[3])
        cx = rect[0] + rect[2]//2
        cy = rect[1] + rect[3]//2
        return (cx, cy)
    else:
        return (w//2, h//2)  # 默认居中

# 处理所有图片
os.makedirs(output_folder, exist_ok=True)
for filename in os.listdir(input_folder):
    if not filename.lower().endswith(('.png', '.jpg', '.jpeg')):
        continue

    img_path = os.path.join(input_folder, filename)
    img = cv2.imread(img_path)
    if img is None:
        continue

    h, w = img.shape[:2]
    target_width = int(target_ratio * h)
    
    if target_width > w:
        print(f"跳过 {filename}：宽度不足")
        continue

    # 获取衣物中心
    cx, _ = detect_clothing_center(img)
    
    # 计算裁剪区域
    x_start = max(0, cx - target_width//2)
    x_end = x_start + target_width
    
    # 边界溢出修正
    if x_end > w:
        x_start = max(0, w - target_width)
        x_end = w

    # 执行裁剪
    cropped = img[0:h, x_start:x_end]
    cv2.imwrite(os.path.join(output_folder, filename), cropped)

print("处理完成！")